import numpy as np
import matplotlib.pyplot as plt
from pathlib import Path
import json
import time
from grid import Grid
from physics import PhysicsModel, PhaseType
from solver import Solver
import matplotlib
# 设置matplotlib显示中文和负数
matplotlib.rcParams['font.sans-serif'] = ['SimHei']  # 设置中文字体为黑体
matplotlib.rcParams['axes.unicode_minus'] = False  # 解决负号显示问题



class SimulationManager:
    """模拟管理器类，负责整个计算过程的控制和数据管理"""
    def __init__(self, config_file: str = 'config.json'):
        # 加载配置
        config_path = Path(__file__).parent.parent / 'config' / config_file
        print(f"尝试加载配置文件: {config_path}")
        
        if not config_path.exists():
            raise FileNotFoundError(f"配置文件不存在: {config_path}")
            
        with open(config_path, 'r') as f:
            self.config = json.load(f)
        
        # 创建结果目录
        self.output_dir = Path(__file__).parent.parent / 'results'
        self.output_dir.mkdir(exist_ok=True)
        
        # 初始化计算网格
        self.grid = Grid(
            n_axial=self.config['grid']['n_axial'],
            channel_length=self.config['grid']['channel_length']
        )
        
        # 初始化物理模型
        self.physics = PhysicsModel()
        
        # 初始化求解器
        self.solver = Solver(self.grid, self.physics)
        
        # 设置求解器参数
        self.solver.max_iterations = self.config['solver']['max_iterations_per_step']
        self.solver.convergence_tolerance = self.config['solver']['convergence_criteria']['pressure']
        self.solver.relaxation_factors = self.config['solver']['relaxation_factors']
        
        # 时间步进参数
        self.current_time = 0.0
        self.total_time = self.config['time_settings']['total_time']
        self.output_interval = 0.1  # 输出间隔
        
        # 性能监控
        self.performance_metrics = {
            'time_steps': 0,
            'total_iterations': 0,
            'computation_time': 0.0
        }
    
    def initialize_flow_field(self):
        """初始化流场"""
        # 设置入口边界条件
        inlet_conditions = self.config['inlet_conditions']
        outlet_conditions = self.config['outlet_conditions']
        
        # 初始化网格和场变量
        self.grid.initialize_field_variables()
        
        # 设置边界条件
        self.grid.set_boundary_conditions(inlet_conditions, outlet_conditions)
        
        print("流场初始化完成")
        print(f"入口压力: {inlet_conditions['pressure']/1e6:.2f} MPa")
        print(f"入口液相速度: {inlet_conditions['liquid']['velocity']:.2f} m/s")
        print(f"入口气相速度: {inlet_conditions['vapor']['velocity']:.2f} m/s")
        print(f"入口空泡份额: {inlet_conditions['alpha']:.3f}")
    
    def save_results(self, time_step: int):
        """保存计算结果"""
        results = {
            'time_step': time_step,
            'time': self.solver.grid.current_time,
            'channels': []
        }
        
        for subchannel in self.solver.grid.subchannels:
            channel_data = {
                'id': subchannel.id,
                'type': subchannel.type.value,
                'geometry': {
                    'area': subchannel.area,
                    'hydraulic_diameter': subchannel.hydraulic_diameter,
                    'wetted_perimeter': subchannel.wetted_perimeter
                },
                'fields': {
                    'pressure': subchannel.pressure.tolist(),
                    'alpha': subchannel.alpha.tolist(),
                    'velocity': {
                        'liquid': {
                            'axial': subchannel.velocity['liquid']['axial'].tolist(),
                            'lateral': subchannel.velocity['liquid']['lateral'].tolist()
                        },
                        'vapor': {
                            'axial': subchannel.velocity['vapor']['axial'].tolist(),
                            'lateral': subchannel.velocity['vapor']['lateral'].tolist()
                        }
                    }
                }
            }
            results['channels'].append(channel_data)
        
        # 保存到文件
        output_file = self.output_dir / f'results_step_{time_step}.json'
        with open(output_file, 'w') as f:
            json.dump(results, f, indent=2)
    
    def plot_results(self, time_step: int):
        """绘制结果图表"""
        fig, axes = plt.subplots(2, 2, figsize=(15, 12))
        
        # 选择一个子通道进行绘制（优先选择中心通道）
        try:
            channel = next(ch for ch in self.grid.subchannels if ch.type.value == 'center')
        except StopIteration:
            channel = self.grid.subchannels[0]
            print(f"警告：未找到中心通道，使用{channel.type.value}通道进行绘制")
        
        z = np.linspace(0, self.grid.channel_length, len(channel.alpha))
        
        # 空泡份额分布
        axes[0,0].plot(z, channel.alpha)
        axes[0,0].set_title('空泡份额分布')
        axes[0,0].set_xlabel('轴向位置 (m)')
        axes[0,0].set_ylabel('空泡份额')
        
        # 压力分布
        axes[0,1].plot(z, channel.pressure/1e6)
        axes[0,1].set_title('压力分布')
        axes[0,1].set_xlabel('轴向位置 (m)')
        axes[0,1].set_ylabel('压力 (MPa)')
        
        # 液相速度分布
        axes[1,0].plot(z, channel.velocity['liquid']['axial'])
        axes[1,0].set_title('液相速度分布')
        axes[1,0].set_xlabel('轴向位置 (m)')
        axes[1,0].set_ylabel('速度 (m/s)')
        
        # 气相速度分布
        axes[1,1].plot(z, channel.velocity['vapor']['axial'])
        axes[1,1].set_title('气相速度分布')
        axes[1,1].set_xlabel('轴向位置 (m)')
        axes[1,1].set_ylabel('速度 (m/s)')
        
        plt.tight_layout()
        plt.savefig(self.output_dir / f'plot_step_{time_step}.png')
        plt.close()
    
    def run_simulation(self):
        """运行模拟计算"""
        print("开始计算...")
        start_time = time.time()
        
        # 初始化
        self.initialize_flow_field()
        time_step = 0
        next_output_time = self.output_interval
        
        # 时间步进循环
        while self.current_time < self.total_time:
            # 求解当前时间步
            self.solver.solve_time_step()
            
            # 更新时间和计数
            self.current_time += self.solver.dt
            time_step += 1
            self.performance_metrics['time_steps'] += 1
            
            # 输出结果
            if self.current_time >= next_output_time:
                print(f"计算时间: {self.current_time:.3f} s")
                self.save_results(time_step)
                self.plot_results(time_step)
                next_output_time += self.output_interval
        
        # 计算性能指标
        self.performance_metrics['computation_time'] = time.time() - start_time
        
        # 输出计算统计信息
        print("\n计算完成!")
        print(f"总时间步数: {self.performance_metrics['time_steps']}")
        print(f"计算耗时: {self.performance_metrics['computation_time']:.2f} s")
        print(f"结果已保存至: {self.output_dir}")

    def solve_time_step(self):
        """求解一个时间步"""
        # 计算自适应时间步长
        dt = self.solver.calculate_time_step()
        
        # 保存当前状态
        for subchannel in self.grid.subchannels:
            subchannel.save_current_state()
        
        # 使用改进的PISO算法求解每个子通道
        for subchannel in self.grid.subchannels:
            # 使用改进的PISO算法
            self.solver.solve_improved_piso(subchannel, dt)
        
        # 更新时间和步数
        self.current_time += dt
        self.current_step += 1
        
        # 检查是否需要保存结果
        if self.current_step % self.save_interval == 0:
            self.save_results()
            self.plot_results()

def main():
    """主程序入口"""
    try:
        # 创建并运行模拟
        simulation = SimulationManager()
        simulation.run_simulation()
        
    except Exception as e:
        print(f"计算过程中出现错误: {str(e)}")
        raise

if __name__ == "__main__":
    main() 